from fastapi import APIRouter, WebSocket, WebSocketDisconnect
import os
import json
import asyncio
import dashscope
import queue
import logging
from dashscope import Generation
from dashscope.audio.tts_v2 import *
from dashscope.audio.asr import *
from logger import logging
from dotenv import load_dotenv
load_dotenv()

router = APIRouter(prefix="/api")

# 初始化API密钥
dashscope.api_key = os.getenv('API_KEY')

prompt = ""

with open('prompt.txt', 'r', encoding='utf-8') as f:
    prompt = f.read().strip()
    logging.info(f"加载提示词: {prompt[:30]}...")


class WSCallback():
    """WebSocket回调类，用于处理TTS流式输出"""

    def __init__(self, websocket: WebSocket, loop: asyncio.AbstractEventLoop):
        self.websocket = websocket
        self.loop = loop
        self.completed = queue.Queue()

    def on_open(self):
        pass

    def on_complete(self):
        self.completed.put(True)

    def on_error(self, message: str):
        logging.error(f'[chat]aliyun.websocket服务异常: {message}')
        self.completed.put(False)

    def on_close(self) -> None:
        pass

    def on_event(self, message) -> None:
        pass

    def on_data(self, data: bytes) -> None:
        def safe_send():
            try:
                task = asyncio.create_task(self.websocket.send_bytes(data))
                task.add_done_callback(self._handle_send_exception)
            except Exception as e:
                logging.error(f'[chat]创建发送任务失败：{e}')
        self.loop.call_soon_threadsafe(safe_send)

    def _handle_send_exception(self, task):
        try:
            task.result()
        except Exception as e:
            logging.error(f'[chat]WebSocket发送数据失败：{e}')


def check_sentence(text):
    """检查字符串是日文、中文还是韩文"""
    word_counts = {
        'japanese': 0,
        'chinese': 0,
        'korean': 0
    }
    for char in text:
        # 检查字符是否在日文相关范围内
        if ('\u3040' <= char <= '\u309F' or  # 平假名
                '\u30A0' <= char <= '\u30FF'):  # 片假名
            word_counts['japanese'] += 1
        elif ('\uAC00' <= char <= '\uD7AF'):  # 韩文
            word_counts['korean'] += 1
        else:
            word_counts['chinese'] += 1
    # 挑选单词数量最大的
    max_key = max(word_counts, key=word_counts.get)
    return max_key


async def llm_task(query: str, websocket: WebSocket):
    """处理LLM对话和TTS任务"""

    try:
        loop = asyncio.get_event_loop()
        synthesizer_callback = WSCallback(websocket=websocket, loop=loop)
        # 创建不同语言的语音合成器
        chinese_english_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('ZH_EN_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        japanese_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('JA_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        korean_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('KO_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        # 调用LLM
        responses = Generation.call(
            model='qwen-plus',
            messages=[
                {'role': 'system', 'content': prompt},
                {'role': 'user', 'content': query}
            ],
            result_format='message',
            stream=True,
            incremental_output=True,
        )
        assistant_response = ''
        cache = ''
        sentence_type = ''
        already_check = False
        for response in responses:
            text_chunk = response.output.choices[0]['message']['content']
            if not text_chunk:
                continue
            assistant_response += text_chunk
            if not already_check:
                cache += text_chunk
                if len(cache) < 10:
                    continue
                else:
                    sentence_type = check_sentence(cache)
                    logging.info(f"检测到的语言类型: {sentence_type}")
                    already_check = True
                    text_chunk = cache
                    cache = ''

            # 发送文本消息
            await websocket.send_text(json.dumps({"text": text_chunk, "type": "text"}))
            # 根据语言类型选择合成器
            if sentence_type == 'japanese':
                japanese_synthesizer.streaming_call(text_chunk)
            elif sentence_type == 'korean':
                korean_synthesizer.streaming_call(text_chunk)
            else:
                chinese_english_synthesizer.streaming_call(text_chunk)
            await asyncio.sleep(0.05)

        # 处理剩余的缓存
        if not already_check and len(cache) > 0:
            sentence_type = check_sentence(assistant_response)
            already_check = True
            logging.info(f"检测到的语言类型: {sentence_type}")
            if sentence_type == 'japanese':
                japanese_synthesizer.streaming_call(cache)
            elif sentence_type == 'korean':
                korean_synthesizer.streaming_call(cache)
            else:
                chinese_english_synthesizer.streaming_call(cache)

        # 完成流式合成
        if already_check:
            if sentence_type == 'japanese':
                japanese_synthesizer.async_streaming_complete()
            elif sentence_type == 'korean':
                korean_synthesizer.async_streaming_complete()
            else:
                chinese_english_synthesizer.async_streaming_complete()

        logging.info(f"AI回复结果:{assistant_response}")

        # 等待TTS完成
        while True:
            if not synthesizer_callback.completed.empty():
                await websocket.send_text(json.dumps({'type': 'end', 'text': assistant_response}))
                break
            await asyncio.sleep(0.1)

    except Exception as e:
        logging.error(f"LLM任务处理失败: {e}")
        await websocket.send_text(json.dumps({'type': 'error', 'message': str(e)}))


async def init_llm_task(query: str, websocket: WebSocket):
    """处理LLM对话和TTS任务"""
    try:
        loop = asyncio.get_event_loop()
        synthesizer_callback = WSCallback(websocket=websocket, loop=loop)
        # 创建不同语言的语音合成器
        chinese_english_synthesizer = SpeechSynthesizer(
            model='cosyvoice-v2',
            voice=os.getenv('ZH_EN_VOICE'),
            format=AudioFormat.PCM_22050HZ_MONO_16BIT,
            callback=synthesizer_callback,
        )
        chinese_english_synthesizer.streaming_call(query)
        for response in query:
            # 发送文本消息
            await websocket.send_text(json.dumps({"text": response, "type": "text"}))
            await asyncio.sleep(0.05)
        chinese_english_synthesizer.async_streaming_complete()
        # 等待TTS完成
        while True:
            if not synthesizer_callback.completed.empty():
                await websocket.send_text(json.dumps({'type': 'end', 'text': query}))
                break
            await asyncio.sleep(0.1)

    except Exception as e:
        logging.error(f"初始化开场白任务处理失败: {e}")
        await websocket.send_text(json.dumps({'type': 'error', 'message': str(e)}))


@router.websocket("/ws/chat")
async def websocket_chat(websocket: WebSocket):
    """WebSocket聊天端点"""
    await websocket.accept()
    logging.info("WebSocket连接已建立")

    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_text()
            logging.info(f"收到客户端消息: {data}")

            try:
                message = json.loads(data)
                if 'text' in message:
                    # 处理聊天消息
                    logging.info(f"接收到的消息: {message['text']}")
                    if message['init']:
                        await init_llm_task(message['text'], websocket)
                    else:
                        await llm_task(message['text'], websocket)
                else:
                    await websocket.send_text(json.dumps({
                        'type': 'error',
                        'message': '消息格式错误，需要包含text字段'
                    }))
            except json.JSONDecodeError:
                await websocket.send_text(json.dumps({
                    'type': 'error',
                    'message': '消息必须是有效的JSON格式'
                }))

    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket错误: {e}")
        try:
            await websocket.send_text(json.dumps({
                'type': 'error',
                'message': f'服务器错误: {str(e)}'
            }))
        except:
            pass
